Repository navigation
JACCL: zeroed send-frame tails, a progress guard, the advertised GID and the upstream JACCL fixes - #33
Merged
Conversation
…rt of ml-explore/mlx ml-explore#4443) Upstream fix. RingImpl::all_gather_wire gave both directions the same slice of the region of each rank, so the upper half of the region of each peer was not transferred. Each direction now has its own half, as in all_reduce_wire. The mesh does not change. Applied without changes, with the upstream test test_all_gather_rank_data. Co-authored-by: Dave Hind <davehind42@gmail.com> Co-authored-by: Anastasiia Filippova <a_filippova@apple.com> (cherry picked from commit 7241f12) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…re/mlx ml-explore#4557) Upstream fix. ReduceScatter::eval_cpu asserted that it has no inputs and then read inputs[0]. A build with assertions stopped at each sum_scatter. The assertion now expects one input. A build without assertions does not change. Applied without changes. (cherry picked from commit 42987b6) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ml-explore#4474) Only the scoped_env helper in python/tests/mlx_tests.py. The test of the next commit (ml-explore/mlx ml-explore#4558) uses it. The upstream commit also changes test_conv.py, test_conv_transpose.py and test_fast_sdpa.py to use the helper. Those changes are not included. (cherry picked from commit ec5f6ff) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…l-explore/mlx ml-explore#4558) Upstream fix. When a backend could not be initialized and strict was false, init registered an empty group under the name of the backend. A later init of the same backend returned that empty group, also when strict was true. A failed initialization is no longer registered, so the later call tries again and a strict call throws. A group that was initialized stays registered until the process exits. This is also true for a JACCL group that a failure closed. Applied without changes, with the upstream test test_strict_init_after_fallback. (cherry picked from commit 8803e52) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
JACCL posts every frame at the full size of its scratch buffer. When a payload was shorter than the buffer, the bytes of earlier sends stayed after it and went to the peer. stage_send_frame() copies the payload and sets the remaining bytes of the frame to zero. All mesh and ring send paths now stage through SharedBuffer::stage_send(). The posted sizes and the receive paths do not change. This is 4e89c2e ("Preserve cluster JACCL send-frame clearing experiment"). It applied without conflicts: the three changed headers are identical in cb77239 and in the parent of that commit. clang-format with the repository style was then applied to the four touched files. Checked on CPU only, with the NativeSendChecks harness (the real JACCL headers, simulated verbs, ASan and UBSan). With the cb77239 headers the stale-tail executable finds old bytes after the payload. With this commit the five check groups pass and the stale-tail executable no longer finds them. No RDMA hardware was used. (cherry picked from commit 4e89c2e) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The mesh and ring polling loops waited for completions without a limit. When a peer died or a completion was lost, the rank used a full CPU core until it was killed. In the incident, a rank that was killed while it waited in MeshImpl::recv left about 160 GB of wired memory with no owning process. The loops also did not read the status of a work completion, so a failed completion was used as data. Every polling loop now calls ProgressGuard::check() after each poll: - A completion with a status other than IBV_WC_SUCCESS is an error. A negative return of ibv_poll_cq is an error, and the shared poll() helpers no longer add it to the count. - When no completion arrives for longer than the limit, the call fails. The limit is JACCL_PROGRESS_TIMEOUT_MS / MLX_JACCL_PROGRESS_TIMEOUT_MS in milliseconds, read once. The default is 120000. Zero or a negative value removes the limit. A value that is not a number keeps the default. In recv the wait includes the time the peer computes, so the default is large. A caller that knows its longest wait can lower it. - The fast path has no clock read. A counter counts the polls without a completion and the clock is read once per 1024 of them. Measured without sanitizers this adds 0.7 ns per poll. A clock read on each poll adds 12.7 ns. The call fails with a std::runtime_error that starts with "[jaccl]" and names the call and the cause. It does not exit the process. MLX runs the collectives on a CPU stream, and Scheduler::enqueue catches the exception and reports it on the thread of the caller, so the caller can do its own cleanup. The errors of the ring pool wires now reach the caller too. Before this change dispatch_wires dropped them because it used future::wait(). State of the group after a failure. It is the same for each exception that leaves a call: - The queue pair and the completion queue of each connection are destroyed, which cancels the posted work, and all buffers are deregistered. The ring does this after all wires returned. The wire that fails sets a flag, and the other wires stop at their next clock sample. - The protection domains and the device contexts stay open until the group is destroyed. MLX keeps the group until the process exits. - Each later call on the group throws immediately. The group cannot be opened again. A new process is necessary. - The peer gets no message. Its own guard ends its wait. The size of the completion queue does not change. Each connection has at most two sends and two receives posted (PIPELINE and NUM_BUFFERS are 2) and its queue has 64 entries. The harness measured a maximum of 4 queued completions. Checked with simulated verbs (the real mesh, ring and rdma headers, ASan, UBSan and TSan; see tests/jaccl in a later commit): a peer that does not answer in each mesh and ring call, a lost completion, a failed completion, a failed poll, calls after a failure, and traffic with waits below the limit. The same cases fail on the parent commit. Checked on two Macs with a Thunderbolt RDMA link and a two-rank mesh, before this commit was moved onto the upstream backports: normal traffic passes with the guard on, so the driver reports a success status on good completions. A rank whose peer ended reported the timeout at the limit and exited, with a model loaded and without one, and neither Mac kept a process or wired memory. The same case without the guard polled for more than nine minutes. A failed completion and the ring path were not run on the real driver. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Connection::info() searches the port's GID table for the IPv4-mapped GID and sends that to the peer as this rank's address. queue_pair_rtr then sets the source GID by a fixed index, 1. The two agree only when the mapped GID happens to sit at index 1. Its position depends on what other addresses the interface carries. On a Thunderbolt port that also has a unique-local IPv6 address the table reads unique-local, link-local, IPv4-mapped, and the mapped GID is at index 2, so the queue pair would send from a link-local IPv6 GID while the peer was told the IPv4-mapped one. Remember the index where the advertised GID was found and use it as the source index. Behaviour is unchanged where the mapped GID is at index 1. On two Macs where the mapped GID is at index 2, normal traffic passes with this commit and also without it, so the Thunderbolt driver tolerates the mismatch there and those runs do not show that the change is necessary. The table shape above was read on an M3 Ultra running macOS 27.2. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Add MLX_JACCL_PROGRESS_TIMEOUT_MS and its JACCL_PROGRESS_TIMEOUT_MS alias to the list of environment variables, with the default and how to remove the limit. Add a section to the JACCL part of the distributed guide that tells what occurs when a peer does not answer: the call raises an error, the group is closed and each later call on it raises an error. The JACCL README now also says that the group cannot be opened again in the same process. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
tests/jaccl is a separate CMake project, in the same way as mlx/distributed/jaccl/lib/examples. The MLX build does not build it. It is not below mlx/, because the install rule of MLX copies each header below mlx/. The tests compile the mesh, ring and rdma headers with a stand-in for <infiniband/verbs.h>. A fake fabric completes a send when it meets a posted receive of the peer. It can lose or fail one completion and it can make the poll fail. - send_frame_tests: each posted frame holds its payload and then only zero bytes, for the mesh and the ring calls. One group does the ring all gather with different data on each rank (ml-explore/mlx ml-explore#4443). - progress_guard_tests: a peer that does not answer in each mesh and ring call, a lost completion, a failed completion, a failed poll, calls after a failure, a refused destroy, waits below the limit, and the values of the environment variable. Each mode is one CTest test because the limit is read once in a process. Run them with ASan and UBSan, or add -DJACCL_TESTS_TSAN=ON for TSan: cmake -S tests/jaccl -B build/jaccl-tests cmake --build build/jaccl-tests ctest --test-dir build/jaccl-tests --output-on-failure The two configurations pass, 8 of 8 tests each. The tests do not use RDMA hardware. No workflow runs them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Add two sections to fork.yaml so that check_forkdiff.py passes. The first one describes the fork work in JACCL: the zeroed send-frame tails, the progress guard, the advertised GID, the documentation of the new environment variable and the tests in tests/jaccl. The second one lists the cherry-picks of ml-explore/mlx ml-explore#4443, ml-explore#4557 and ml-explore#4558 and the scoped_env helper of ml-explore#4474 as upstream work. It also says why ml-explore#4607 is not ported. base.hash does not change, because a cherry-pick does not move the merge-base. The overview now also names the JACCL changes, which are not gated, and counts three sections that are not fork work. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The Check Lint job runs cmake-format on all files and wants one argument per line in this foreach. No test changes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
The slice bounds that came with the upstream all_gather fix multiply lr and n_wires as int before the byte count widens the product. Both are small, so nothing overflows today, but the code scanner reports the pattern as a high-severity alert on the pull request. Cast the first factor so the whole product is int64_t. No behaviour change; the simulated-verbs tests pass. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Upstream ml-explore/mlx c7ff35d ("Skip unnecessary simdgroup computations for quantised MOE matmuls on NAX", ml-explore#4352) reached this fork through the v0.32.2 squash (734241b), not by ancestry, in the form that ml-explore#4572 later gave it: `const bool sg_active = sgp_sm > 0;` and two `if (sg_active)` blocks around the MMA loop in each of fp_gather_qmm_rhs_nax (fp_quantized_nax.h) and affine_gather_qmm_rhs_nax (quantized_nax.h). This removes the flag and runs the MMA in every simdgroup again, as before ml-explore#4352. It is the same change as the earlier revert of ml-explore#4352 on 0.32.2 (e9debbe) and its re-application on the 0.32.3 kernels, hand-ported because the surrounding lines differ here. Why: these kernels run only on GPUs with NAX. On a two-Mac JACCL pair with one NAX GPU, mixture-of-experts serving with the skip present lost all-reduce completions, and 92 to 96 GiB stayed wired on the host until it was rebooted. The builds that served that pair afterwards carried the revert. It is therefore kept out of every build that runs distributed on mixed hardware. Output is unchanged: stores stay bounded by sgp_sm, and the rows past the tile that the MMA now reads are either inside M (rows_in_bounds) or loaded as zeros by load_safe with sgp_sm rows. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ion queues A JACCL call that failed (no completion for the progress limit, a failed work completion, a failed poll) closed its group and threw. The thrown error reaches the program only through the stream that ran the collective, and a rank whose peer died sat in the guard for two minutes before it could even start to release its memory. RDMA work left outstanding on a dying process has been seen to keep tens of GiB wired until a reboot. - The default progress limit is 30000 ms (was 120000). - After the group is released (queue pairs and completion queues destroyed, buffers deregistered; the ring waits for all its wires), the default action is now a teardown exit: print the reason, run the memory release hook, exit with status 75 (EX_TEMPFAIL, restartable) through std::_Exit. JACCL_TIMEOUT_ACTION / MLX_JACCL_TIMEOUT_ACTION = throw keeps the previous behaviour; any other value keeps the default. Only the first failing thread exits; a second one waits for it. - MLX's jaccl.cpp registers the hook when it creates a group: set_wired_limit(0), then clear_cache(), each best effort. The hook gets 10 s, then the process exits anyway. - Completion queues have 4 * (MAX_SEND_WR + MAX_RECV_WR) entries, four times what a queue pair can have outstanding, so that a provider never sees a full queue. The queue pair attributes are not touched (the driver refuses changes to the RTR mask) and no global collective lock is added. Tests (tests/jaccl, simulated verbs, ASan/UBSan and TSan, all pass): the existing modes select the throw action and are otherwise unchanged. New teardown-exit mode runs each case as a child process and checks status 75, the reason, one release after the group was released and the time from the start of the call to the exit: mesh recv and all_reduce, ring recv on two wires, a failed completion (at once), no hook, a hook that never returns (cut off at 10 s), two groups failing together (one release, one exit), an unknown action, and the throw action by name and by alias. default-exit runs a stalled call with no variable set and checks the exit at 30 s. send_frame reports the completion high-water mark against the new depth. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jonathan308
added a commit
to Layr-Labs/mlx-swift
that referenced
this pull request
Oct 10, 2026
Moves Source/Cmlx/mlx from d0f8b0a8 to fdae3174 on Layr-Labs/mlx darkbloom/jaccl-send-frame-progress-guard (draft Layr-Labs/mlx#33): - b7658964: ring all_gather slice arithmetic in 64 bits. - ccf517a0: revert of the NAX quantised-MoE gather_qmm simdgroup skip (upstream ml-explore/mlx c7ff35d97), which lost JACCL all-reduce completions on a two-Mac pair with one NAX GPU. Output unchanged. - fdae3174: 30 s progress limit, teardown exit with a memory release hook (status 75) by default, completion queues sized 4x. No source files were added or removed, so the package manifest is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… section fdae317 changed three files that no fork.yaml glob covered, so the "fork.yaml describes the fork" check failed on this branch: mlx/distributed/jaccl/jaccl.cpp (registers the memory release hook that the teardown exit runs) and lib/jaccl/mesh.cpp / ring.cpp (create each completion queue at COMPLETION_QUEUE_DEPTH). Add the three paths to the JACCL section's globs and name them in its description. No code change. scripts/check_forkdiff.py --upstream-ref upstream/main: 205 changed files, 205 described, exit 0 (was 3 uncovered, exit 1). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This was referenced Oct 15, 2026
Brings in main at f78f0fa (2026-10-24): the upstream ml-explore/mlx sync to afc03c5 (#38, with its history restored by #42) and the fork fixes #40 and #43-#46. Conflicts and how they were resolved: - mlx/distributed/jaccl/lib/jaccl/mesh.cpp: main's per-wire queue setup (upstream ml-explore#4607, a second cable for the mesh) with this branch's COMPLETION_QUEUE_DEPTH for every completion queue. - mlx/distributed/jaccl/lib/jaccl/mesh_impl.h: main's stride parameters for all_gather and sum_scatter, with this branch's progress-guard calls and stage_send in every send path. - mlx/distributed/jaccl/lib/jaccl/ring_impl.h: this branch keeps its guarded member dispatch_wires (pool-wire errors are reported, the group closes after all wires return). Main moved an unguarded copy into threadpool.h and called it as dispatch_wires(pool_, ...); the ring call sites call the member again. The all_gather slice code is main's (identical to this branch's backport of upstream ml-explore#4443). - quantized_nax.h, fp_quantized_nax.h: main's kernels. Main already keeps the revert of the upstream ml-explore#4352 simdgroup skip (no sg_active) and adds the ragged-K tail bound of upstream ml-explore#4009. - fork.yaml: main's file, plus the JACCL description and section. The upstream backports this branch carried (ml-explore#4443, ml-explore#4557, ml-explore#4558, the scoped_env helper of ml-explore#4474) are now upstream ancestry, so their "not fork work" section and its references are gone. The JACCL section now says that a mesh with more than one wire has one guard per wire. Also: - python/tests/ring_test_distributed.py: main's file. The auto-merge had added test_strict_init_after_fallback twice. - tests/jaccl: the mesh all_gather and sum_scatter calls pass main's new stride argument (equal to the chunk size, as main's own all_reduce_scatter_gather does). Tests: tests/jaccl (simulated verbs, ASan and UBSan): 10 of 10 passed. check_forkdiff.py --upstream-ref upstream/main: 111 changed files, all described, base afc03c5 is the merge-base. clang-format 21: no changes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The merge of main brought upstream ml-explore#4607 (a second cable for the JACCL mesh). Its MeshGroup::split_wires runs the wires through the free dispatch_wires in threadpool.h, which waits on the pool futures but never calls get(), so an error thrown in a pool wire was dropped and the call returned with that wire's part of the output unwritten. This branch had fixed the same drop in the ring's dispatch_wires; this applies the fix to the copy that main moved into threadpool.h. After every wire returned, the first pool-wire error is thrown. Nothing changes when no wire fails or when there is one wire (the two-Mac cluster uses one). fork.yaml: the JACCL section describes it and lists threadpool.h. Tests: tests/jaccl progress_guard_tests gains "dispatch_wires throws the error of a pool wire" (fails without the fix, passes with it). tests/jaccl with ASan and UBSan: 10 of 10 passed; with TSan: 10 of 10 passed. check_forkdiff.py --upstream-ref upstream/main: 112 changed files, all described. clang-format 21: no changes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jonathan308
added a commit
to Layr-Labs/mlx-swift
that referenced
this pull request
Oct 26, 2026
Restitches the nested pins after the merge of main 6959486: - Source/Cmlx/mlx: fdae3174 -> 187b6c67. That is the merge of core main f78f0fa9 into Layr-Labs/mlx#33 (the JACCL send-frame, progress-guard and GID fixes), plus its pool-wire error report. - Source/Cmlx/mlx-c: 489e9654 (mlx-c#14) -> 004c735f, main's pin. The merged C layer 73720efb (main into mlx-c#15) cannot be used here: it calls core functions that only mlx#35 has (metal_kernel_with_fenced_inputs and the sealed collectives), so it does not compile against this core. mlx-c#14 itself predates the C API that main's Swift calls, and it is not merged anywhere. Regenerated as tools/update-mlx.sh does from those pins: - The two mirrored distributed_group.h copies drop mlx_distributed_init_jaccl_with_bootstrap, which 004c735f does not implement. The bootstrap bridge comes back with mlx-c#15 in the next layer. - Every other mirrored header, the embedded JIT and CPU sources and the Xcode project regenerate byte-identical. The include-framework Metal.hpp differs only in its "Autogenerated" line and is kept as it is. fork.yaml's pin paragraph now names this layer's core and C. Tests (Ultra A, Xcode 27.1, at these pins): - tools/update-mlx.sh regeneration from the two pins matches the tree. - check-generated-metal-sources.py: pass. - swift build --build-tests: complete. - Fork Tests, all 3 shards of the CI plan, with an AOT mlx.metallib from this core: - XCTest: 852 listed, 848 passed, and the 2 permitted skips. - The fresh-process allocation-cap test and both multi-process ring suites passed. - Swift Testing: 749 of 749 passed. - test-shards.py verify: every test ran exactly once. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…nto the JACCL guard branch Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jonathan308
marked this pull request as ready for review
October 27, 2026 18:49
main's #48 changed mlx/backend/metal/jit_kernels.cpp without a fork.yaml section, so the fork.yaml gate failed on main and on every PR based on it. The new section explains the change: the JIT source passes attention_nax_dsplit the same template arguments as its metallib instantiation. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
davidtai
added a commit
to Layr-Labs/mlx-swift
that referenced
this pull request
Oct 27, 2026
build: pin mlx to the JACCL fixes (Layr-Labs/mlx#33); keep main's mlx-c
…fsets; report refused deregistration - The wires of a MeshGroup now share one failure flag. A failed wire stops a call that waits on another wire (WireStopped), and every later call on the group is refused, as the README and the guard comment say. Before, each MeshImpl had its own flag, so a failure on wire 1 left wire 0 open. - MeshImpl::all_gather keeps read_offset and write_offset in int64_t. With int, more than 2 GiB per rank overflowed and indexed out of bounds. - SharedBuffer::deregister reports a region that the verbs refuse to deregister and keeps it for the destructor, as release() does for queue pairs and completion queues. Before, it dropped the handle silently. - tests/jaccl: a refused-deregister case and a two-wire mesh case with a shared flag. Fork Tests build-cpp now builds and runs tests/jaccl, which no CI job built before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Review fixes in 10d25ed:
Not changed (design question): the default 30 s limit plus teardown-exit also ends a healthy rank that waits more than 30 s for a peer (pipeline recv during a long prefill, idle server collectives). It is documented, and callers can raise the limit with JACCL_PROGRESS_TIMEOUT_MS. |
davidtai
added a commit
that referenced
this pull request
Oct 27, 2026
…d all-sum branch mesh_impl.h keeps this branch's per-wire StartupWaitBudget and takes #33's shared failure flag (std::shared_ptr), so a failed wire still closes every wire of the group. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
davidtai
added a commit
that referenced
this pull request
Oct 27, 2026
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
davidtai
added a commit
that referenced
this pull request
Oct 27, 2026
…ering branch Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
davidtai
added a commit
to Layr-Labs/d-inference
that referenced
this pull request
Oct 27, 2026
The closure carries the runtime fixes of the MLX review: the JACCL mesh failure flag and 64-bit all_gather offsets (Layr-Labs/mlx#33), fenced inputs kept on device and sealed frames over INT_MAX refused (Layr-Labs/mlx#35), the fast fence before the event fence (Layr-Labs/mlx#41), and GLM5NextClusterImagePlan.inspectPNG offsets taken from the Data's startIndex (Layr-Labs/mlx-swift-lm#308). Pins: libs/mlx-swift-lm d5bd32b8, libs/mlx-swift c633d221 (the revision mlx-swift-lm's Package.swift pins), libs/mlx 72dffa2d (mlx-swift's Source/Cmlx/mlx). mlx-c in that mlx-swift is 2ac897e3. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Three fixes to JACCL, the RDMA-over-Thunderbolt backend, on current
main, with the upstream JACCL fixes that the fork did not have in front of them.This PR contains and extends #26. #26 carries the send-frame commit alone on an older base. Here the same commit (authorship kept) is on current
main, formatted with clang-format and described infork.yaml, which were the two open items of #26. #26 can be closed when this one is accepted.Status: draft. The mesh path was run on two Macs over Thunderbolt RDMA (results below). The ring path was run only with simulated verbs.
Commits
d92347bad3739307461496887112cbf895e89385cbe63d148785a85ec9dcfb50fc71f7fbd0f8b0a84241bc45b7658964Upstream backports
Upstream
mainwas read at0d8236fa(2026-10-09). Five upstream commits after the fork base (77a0c1e8) touchmlx/distributed/, and one more touches the CPU code of a distributed op.ml-explore/mlx#4210adf21deaml-explore/mlx#4281a44fc8c0ml-explore/mlx#44437241f12eall_gather: both directions sent the same slice, so the upper half of each peer's region was not transferredb7658964) casts the first factor of its slice arithmetic to 64 bits, because the code scanner flags the int productml-explore/mlx#455742987b66ReduceScatter::eval_cpuasserted that it has no inputs, so a build with assertions stopped at eachsum_scatterml-explore/mlx#45588803e524initreturned itml-explore/mlx#4607f4232c7fscoped_envtest helper ofml-explore/mlx#4474. Only that helper is taken from Add a scoped_env utility for tests ml-explore/mlx#4474. The other changes of Add a scoped_env utility for tests ml-explore/mlx#4474 (test refactors) are not included.MeshImplfor each wire and movesdispatch_wiresout ofRingImpl. The progress guard of this PR closes oneMeshImpland joins the ring wires insideRingImpl::dispatch_wires, so the guard must be reworked for that structure and tested again on hardware. That is a separate PR.0d8236fathe JACCL loops do not readwc.status, have no time limit, post full frames without clearing, and usesgid_index = 1.The three fixes
1. Send-frame tail (
95e89385)JACCL posts each frame at the full size of its scratch buffer. When the payload is shorter than the buffer, the bytes of earlier sends stay after it and go to the peer.
stage_send_frame(newsend_frame.h) copies the payload and sets the rest of the frame to zero. All mesh and ring send paths stage throughSharedBuffer::stage_send. The posted sizes and the receive paths do not change.2. Progress guard (
cbe63d14)Before: each mesh and ring polling loop waited for completions without a limit and did not read the status of a work completion. A rank whose peer had stopped used a full core until it was killed. A failed completion was used as data.
After:
ProgressGuard::check()(newprogress_guard.h) runs after each poll of each loop (6 loops inmesh_impl.h, 4 inring_impl.h).IBV_WC_SUCCESSis an error at once. A negative return ofibv_poll_cqis an error at once.std::runtime_errorthat starts with[jaccl]and names the call and the cause. It does not exit the process.Scheduler::enqueuecatches the exception on the CPU stream and reports it to the caller, so the caller can do its own cleanup.dispatch_wiresdropped them before, because it usedfuture::wait().New environment variable.
JACCL_PROGRESS_TIMEOUT_MS, with the aliasMLX_JACCL_PROGRESS_TIMEOUT_MS, is the longest time in milliseconds that a call waits for a completion. The default is 120000.0or a negative value removes the limit. A value that is not a number keeps the default. The value is read once in a process. The wait of a rank includes the time that its peer computes (for example inrecv), so the limit must be above the longest such wait.The group is closed for good after a failure. This is the same for each exception that leaves a call:
3. Advertised GID (
8785a85e)Connection::info()finds the IPv4-mapped GID of the port and sends it to the peer.queue_pair_rtrthen set the source GID by the fixed index 1. The two agree only when the mapped GID is at index 1. The connection now keeps the index where it found the advertised GID and sends from that index. Nothing changes where that index is 1.On the two test Macs the mapped GID is at index 2. The commit passed normal traffic on that pair. The code without it also passes on that pair, so the hardware runs do not show that the fix is necessary.
Tests with simulated verbs
tests/jaccl(new) compiles the realmesh_impl.h,ring_impl.handrdma.hwith a stand-in for<infiniband/verbs.h>. A fake fabric completes a send when it meets a posted receive of the peer. It can lose or fail one completion and it can make the poll fail. The directory is a separate CMake project, in the same way asmlx/distributed/jaccl/lib/examples, so the MLX build does not build it and no header of it is installed.cmake -S tests/jaccl -B build/jaccl-tests # add -DJACCL_TESTS_TSAN=ON for TSan cmake --build build/jaccl-tests ctest --test-dir build/jaccl-tests --output-on-failureAt
d0f8b0a8: 8 of 8 CTest tests pass with ASan and UBSan, and 8 of 8 pass with TSan. At the head (4241bc45, the format change): 8 of 8 pass with ASan and UBSan.The same test sources were also built against the headers of each commit of this branch (a private harness that exports one revision at a time):
main(cb77239b)all_gatherdefectd92347bato7112cbf8)all_gather), 5 fail95e89385)cbe63d14) and laterThe 34 guard cases that fail without the guard:
The 4 cases that pass without the guard are three transfers with waits below the limit and the count of the fake objects. With the guard, each stalled call fails between its limit (250 ms in the tests) and one second after it, the queue pairs and completion queues of the rank are destroyed exactly once, each region is deregistered exactly once, and each later call is refused in less than 100 ms. With the limit removed, a stalled call continues to wait and a failed completion is still an error.
Runs on two Macs
Two Apple silicon Macs (an M3 Ultra and an M5 Max, macOS 27) with a Thunderbolt RDMA link and a two-rank JACCL mesh. These runs used the three fix commits before the rebase onto the backports. The rebase changes the ring
all_gathercode, one assertion and the registration of a failed initialization. None of these is on the mesh path that was run.JACCL_PROGRESS_TIMEOUT_MS=4000[jaccl] mesh all_reduce: no completion for 4001 ms … The group is closedand exited about 4 s after its peer ended.mesh recv: no completion for 60001 msand exited by itself. Neither Mac kept a process or wired memory. The next session gave exact results.jaccl::MeshImpl::all_reduce→tbt_poll_cqat 100% CPU for more than nine minutes.CI on this pull request
At
4241bc45:metal-build-and-test(macOS, deployment target 26.2): the JACCL library (tcp,rdma,mesh,ring,jaccl) andmlx/distributed/jaccl/jaccl.cppare compiled and linked. 292 C++ test cases pass and 1,869 Python tests run without a failure.No CI job runs JACCL traffic: the runners have no RDMA device.
Not shown
[jaccl] Could not destroy a queue pair (N)and keeps the handle for its destructor.test_all_gather_rank_data,test_strict_init_after_fallback) were not run locally.ring_test_distributed.pyin the Fork Tests workflow runs them; that workflow did not run on this pull request.tests/jacclin CI. No workflow builds it.Fork bookkeeping
fork.yamlhas a section for the JACCL fork work and a section that lists the cherry-picks as upstream work.python3 scripts/check_forkdiff.py --upstream-ref refs/remotes/upstream/mainpasses: 202 changed files, 202 described.base.hashdoes not change, because a cherry-pick does not move the merge-base.foreachoftests/jaccl/CMakeLists.txt, and4241bc45applies that. The job passes at the head.docs/src/usage/environment_variables.rstanddocs/src/usage/distributed.rstdescribe the new variable and what occurs when a peer does not answer.How AI was used
Claude Code (Anthropic) wrote the two fork fixes after the send-frame commit, the tests, the documentation changes, the backport selection, the commit messages and this description, and ran the simulated-verbs tests and the two-Mac sessions, at the direction of the account owner, who asked for this draft. The send-frame commit is by Gajesh Naik and is unchanged apart from formatting. The upstream commits are cherry-picks with their original authors.
AGENTS.mdandCLAUDE.mdin this repository are upstream's files and say that an agent must not push or open a pull request; this draft was opened on the explicit request of the account owner, in this fork only. Nothing is proposed toml-explore/mlx.Update 2026-10-10: guard v2 and the NAX MoE simdgroup-skip revert
Two commits on top of
b7658964:ccf517a0reverts upstreamc7ff35d97(the NAX quantised-MoE simdgroup skip, ml-explore/mlx PR 4352), which reached this fork through the v0.32.2 squash (734241bb), not by ancestry. On a two-Mac JACCL pair with one NAX GPU, MoE serving with the skip lost all-reduce completions and left 92 to 96 GiB wired until reboot. Output is unchanged (stores stay bounded bysgp_sm).fdae3174: default progress limit 30 s (was 120 s); after the group is released the default action is a teardown exit (memory release hook:set_wired_limit(0)thenclear_cache(), 10 s budget, then exit status 75);JACCL_TIMEOUT_ACTION=throwkeeps the old behaviour; completion queues sized 4x the outstanding work. Queue pair attributes untouched; no global collective lock.Checked:
ccf517a0; ASan/UBSan 10/10 and TSan 10/10 atfdae3174(new teardown-exit and 30 s default-exit cases).exactagainst the pair runs before the revert, speeds within run-to-run spread, 56 of 56 sustained requests each with no lost completion; on the M5 alone the revert is bit-identical in output and speed.Note for consumers: the d-inference
libs/mlxpin (cb77239b) and the MLX nested in itslibs/mlx-swiftalso carry the skip today; they need to move to this head (with a rebuilt metallib) for any distributed build.🤖 Generated with Claude Code
Progress 2026-10-26: main merged, stack re-pinned, validated on two M5 Ultra 256 GB
Head
187b6c67: main merged (ed96e1ab) plus the multi-cable dispatch_wires error fix (187b6c67).MLX-MERGE-2026-10-26.md(core),MLX-C-MERGE-2026-10-26.md,MLX-SWIFT-MERGE-2026-10-26.md,MLX-SWIFT-LM-MERGE-2026-10-26.md(kept with the overnight records).187b6c67, test fails without the fix). Cluster collectives, quantized numerical fixes and optional Metal allocation cap #35 carries Fix Metal fence ordering and sealed collective portability #41's barrier line so it does not undo main's [BUG] Metal fence deadlock ml-explore/mlx#4689. Main's value-based fence waits are taken; the cluster keeps its single spin ordering. Upstream JIT bug fixed (ce21a37d): head-dim-256 NAX attention failed to build its run-time template on M5. Fix Metal fence ordering and sealed collective portability #41 addsMLX_QMV_BATCH_INVARIANT(caa9ed53, off by default): quantized products of 2 to the vector limit - 1 rows stay on qmv_wide, so a row's bits do not depend on how many rows share its launch (main's [Metal] Add qmv_mma/qmv_nax kernels for 6-16 rows with large N/K ml-explore/mlx#4712/[Metal] Add staged affine_qmm_t_nax kernels with split-k ml-explore/mlx#4682 routes otherwise change them). Cluster processes turn it on; single-Mac keeps main's kernels.Source/Cmlx/mlx= the matching core layer,Source/Cmlx/mlx-c= mlx-c Expose coherent allocator accounting and correct quantized bias sums #15, Docker: xcrun: not found ml-explore/mlx#54 on main's C); generated sources regenerated withtools/update-mlx.shand compared; random.uniform with dtype float16 returns all 0s ml-explore/mlx#61 pins corecaa9ed53(onlyinclude-framework/mlx-utils.hchanges). Fork Tests (all three shards with an AOT metallib): XCTest 849/853 + 2 permitted skips, Swift Testing 749/749.Nothing merged or released; the owners sign and release.
Update 2026-10-26 17:15 PDT: main/master 0.10.3 merged through the stack
Head
283017ce: mlx main6bc9de5bmerged (#48, the head-dim split NAX attention JIT fix behind provider 0.10.3). Every PR in the cluster stack (mlx #33 → #35 → #41, mlx-c #15, mlx-swift ml-explore#54 → ml-explore#56 → ml-explore#61, SDK ml-explore#308 → ml-explore#352, d-inference ml-explore#1407 → ml-explore#1426 / ml-explore#1577, platform ml-explore#54) is again mergeable in that order, as drafts.Review fixes and follow-up (2026-10-27)
Head
28e32c79: review fixes in10d25ed6(shared mesh failure flag across wires, 64-bitall_gatheroffsets, refused deregistration reported, tests/jaccl built and run in Fork Tests) and the clang-format fix28e32c79.Follow-up: the 30 s default of the progress guard applies only to callers that do not set
JACCL_PROGRESS_TIMEOUT_MS. A healthy rank that waits for a peer for longer than the limit (a pipeline recv during a long prefill, idle server collectives) is ended, so such callers must set the limit. The Darkbloom cluster worker sets it for both ranks from--progress-timeout-ms(default 60000) and refuses a guarded run without it.